from typing import Optional

from IPython.core.display import Image
from IPython.core.display_functions import display
from langgraph.graph import StateGraph
from langgraph.graph import START, END
from pydantic import BaseModel


class MyState(BaseModel):
    x: int
    result: Optional[str] = None


def check_x(state: MyState)->MyState:
    print(f"[check_x]Received state: {state}")
    return state

def is_even(state: MyState) -> bool:
    print(f"[is_even]Received state: {state}")
    return state.x % 2 == 0

def handle_even(state: MyState) -> MyState:
    print(f"[handle_even]Received state: {state} 是偶数")
    return MyState(x=state.x , result="Even")

def handle_odd(state: MyState) -> MyState:
    print(f"[handle_odd]Received state: {state} 是奇数")
    return MyState(x=state.x , result="Odd")

if __name__ == '__main__':
    # pip install pyppeteer ipython
    builder = StateGraph(MyState)
    builder.add_node("check_x", check_x)
    builder.add_node("handle_even", handle_even)
    builder.add_node("handle_odd", handle_odd)
    #添加分支条件
    builder.add_conditional_edges(
        "check_x",
        is_even,
        {
            True: "handle_even",
            False: "handle_odd",
        },
    )
    builder.add_edge(START, "check_x")
    builder.add_edge("handle_even", END)
    builder.add_edge("handle_odd", END)
    print(builder.edges)
    print(builder.nodes)
    graph = builder.compile()
    graph_image = graph.get_graph(xray=True).draw_mermaid_png()
    with open("graph_conditional.png", "wb") as f:
        f.write(graph_image)
    print("Graph saved to graph.png")
    print(f"测试5是否是偶数: {graph.invoke(MyState(x=5))}")
